from easymoney.framework.mysql.core.db_factory import InvestmentRecommendationEntity, db
from peewee import IntegrityError


def save(code, name, price, pe_ttm):
    entity = InvestmentRecommendationEntity(
        stock_code=code,
        stock_name=name,
        stock_price=price,
        pe_ttm = pe_ttm
    )
    # 保存到数据库
    with db:
        # 在这里执行数据库操作
        try:
            entity.save()
        except IntegrityError as e:
            error_message = str(e)
            parts = error_message.strip("()").split(", ")
            raise IntegrityError(parts[0], parts[1])
        except Exception as e:
            print(e)


def update_frequency(code, price, pe_ttm):
    entity = InvestmentRecommendationEntity.get(InvestmentRecommendationEntity.stock_code == code)
    entity.frequency = entity.frequency + 1
    entity.stock_price = price
    entity.pe_ttm = pe_ttm
    entity.save()


def update_price(code, price):
    entity = InvestmentRecommendationEntity.get(InvestmentRecommendationEntity.stock_code == code)
    entity.stock_price = price
    entity.save()


def exists(code):
    return InvestmentRecommendationEntity.select().where(InvestmentRecommendationEntity.stock_code == code).exists()


def get_all():
    """
    查询 InvestmentRecommendationEntity 表的所有数据
    :return: 包含所有记录的列表（列表中的每个元素是 InvestmentRecommendationEntity 模型实例）
    """
    try:
        # 执行查询并返回所有记录
        return list(InvestmentRecommendationEntity.select())
    except Exception as e:
        print(f"查询数据失败: {str(e)}")
        return []
